module Tools

using LinearAlgebra
using ..FLEXINVERT: R_earth, π

# Port of mod_tools: sort (returning indices), unique, grid_area, interp, average

"""
    argsort_indices(x::AbstractVector{<:Real}) -> inds::Vector{Int}

Returns the permutation indices that sort the vector in ascending order.
Fortran's routine also returned the sorted array; in Julia you can use
`x[argsort_indices(x)]` if needed.
"""
function argsort_indices(x::AbstractVector{<:Real})
    return sortperm(x)
end

"""
    unique_stable(x::AbstractVector{T}) where T -> Vector{T}

Stable unique preserving first occurrence order.
"""
function unique_stable(x::AbstractVector)
    seen = Set{eltype(x)}()
    out = eltype(x)[]
    for v in x
        if !(v in seen)
            push!(seen, v); push!(out, v)
        end
    end
    return out
end

"""
    grid_area(ylat::Real, dy::Real, dx::Real)

Area of a lat-lon grid cell centered at `ylat` with extents `dy` and `dx` (degrees).
Matches Fortran's spherical approximation.
"""
function grid_area(ylat::Real, dy::Real, dx::Real)
    rad = π/180
    ylata = ylat
    ylatp = ylat + 0.5*dy
    ylatm = ylat - 0.5*dy
    hzone = if ylatm < 0 && ylatp > 0
        R_earth * rad
    else
        cosfact  = cos(ylata*rad) * R_earth
        cosfactp = cos(ylatp*rad) * R_earth
        cosfactm = cos(ylatm*rad) * R_earth
        if cosfactp < cosfactm
            sqrt(R_earth^2 - cosfactp^2) - sqrt(R_earth^2 - cosfactm^2)
        else
            sqrt(R_earth^2 - cosfactm^2) - sqrt(R_earth^2 - cosfactp^2)
        end
    end
    return 2*π*R_earth*hzone*dx/360
end

"""
    interp(datain, timein, timeout) -> dataout

Linear interpolation of datain (nx × nin) from `timein` to `timeout` (both monotonic).
"""
function interp(datain::AbstractMatrix{<:Real}, timein::AbstractVector{<:Real}, timeout::AbstractVector{<:Real})
    nx, nin = size(datain)
    nout = length(timeout)
    dataout = zeros(eltype(datain), nx, nout)
    for n in 1:nout
        # find nearest time index
        ii = argmin(abs.(timein .- timeout[n]))
        i1 = timein[ii] > timeout[n] ? max(1, ii-1) : ii
        i2 = timein[ii] > timeout[n] ? ii : min(nin, ii+1)
        if i1 != i2
            w = (timeout[n] - timein[i1]) / (timein[i2] - timein[i1])
            @inbounds dataout[:, n] .= datain[:, i1] .+ (datain[:, i2] .- datain[:, i1]) .* w
        else
            @inbounds dataout[:, n] .= datain[:, i1]
        end
    end
    return dataout
end

"""
    average(datain, timein, timeout) -> dataout

Average higher-frequency datain (nx × nin) defined on `timein` onto coarser
`timeout` bins, preserving means as in the Fortran implementation.
"""
function average(datain::AbstractMatrix{<:Real}, timein::AbstractVector{<:Real}, timeout::AbstractVector{<:Real})
    nx, nin = size(datain)
    nout = length(timeout)
    fin = 1.0 / (timein[2] - timein[1])
    fout = 1.0 / (timeout[2] - timeout[1])
    fin < fout && error("average: input series less frequent than output")
    dataout = zeros(eltype(datain), nx, nout)
    i0 = argmin(abs.(timein .- timeout[1]))
    frac = fin / fout
    for n in 1:nout
        i1 = max(1, i0 + Int(floor((n-1)*frac)))
        i2 = min(nin, i0 + Int(floor(n*frac)) - 1)
        rem1 = (n-1)*frac - floor((n-1)*frac)
        rem2 = n*frac - floor(n*frac)
        if rem1 > 0
            if i2 < nin
                dataout[:, n] .= (sum(datain[:, (i1+1):i2]; dims=2) .+ rem1 .* datain[:, i1] .+ rem2 .* datain[:, i2+1]) ./ (i2 - i1 + rem1 + rem2)
            else
                dataout[:, n] .= (sum(datain[:, (i1+1):i2]; dims=2) .+ rem1 .* datain[:, i1]) ./ (i2 - i1 + rem1)
            end
        else
            if i2 < nin
                dataout[:, n] .= (sum(datain[:, i1:i2]; dims=2) .+ rem2 .* datain[:, i2+1]) ./ (i2 - i1 + 1 + rem2)
            else
                dataout[:, n] .= sum(datain[:, i1:i2]; dims=2) ./ (i2 - i1 + 1)
            end
        end
    end
    return dropdims(dataout; dims=2) isa AbstractVector ? reshape(dataout, nx, nout) : dataout
end

end # module

